package com.ask.forMe.langchain4j.rag.splitter;

import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.document.splitter.HierarchicalDocumentSplitter;

/**
 * 自定义的文档分割器，按标题分割1
 */
public class DocumentByHeadingSplitter extends HierarchicalDocumentSplitter {
    private final int minHeadingLevel;
    private final int maxHeadingLevel;

    public DocumentByHeadingSplitter(int maxSegmentSizeInChars, int maxOverlapSizeInChars,
                                     int minHeadingLevel, int maxHeadingLevel) {
        super(maxSegmentSizeInChars, maxOverlapSizeInChars, null, null);
        this.minHeadingLevel = minHeadingLevel;
        this.maxHeadingLevel = maxHeadingLevel;
    }

    @Override
    protected String[] split(String text) {
        // 使用正则表达式匹配 Markdown 格式的标题
        String headingPattern = RegexSplitUtils.generateHeadingPattern(minHeadingLevel, maxHeadingLevel);
        String[] split = text.split(headingPattern);
        return split;
    }

    @Override
    protected String joinDelimiter() {
        return "\n\n";
    }

    @Override
    protected DocumentSplitter defaultSubSplitter() {
        return new DocumentByParagraphSplitter(this.maxSegmentSize, this.maxOverlapSize, this.tokenizer);
    }
}
